package com.walter.springai.config;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import jakarta.annotation.Resource;
import org.jetbrains.annotations.NotNull;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Lazy;
import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.serializer.RedisSerializer;
import org.springframework.data.redis.serializer.SerializationException;
import org.springframework.data.redis.serializer.StringRedisSerializer;
import org.springframework.stereotype.Component;

import java.io.ByteArrayOutputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 基于R edis 实现历史会话管理
 */
@Component
public class RedisChatMemory implements ChatMemory {

    @Resource
    @Lazy
    private RedisTemplate<String, List<Message>> redisChatTemplate;

    @Override
    public void add(String conversationId, List<Message> messages) {
        getAndSetAll(conversationId, messages);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> all = getAndSetAll(conversationId, null);
        return all.stream().skip(Math.max(0, all.size() - lastN)).toList();
    }

    private List<Message> getAndSetAll(String conversationId, List<Message> newMessages) {
        String key = getConversationKey(conversationId);
        List<Message> messages = redisChatTemplate.opsForValue().get(key);
        if (messages == null) {
            return List.of();
        }
        if (newMessages != null) {
            messages.addAll(newMessages);
            redisChatTemplate.opsForValue().set(key, messages);
            redisChatTemplate.expire(key, 1, TimeUnit.DAYS);
        }
        return messages;
    }

    @Override
    public void clear(String conversationId) {
        String key = getConversationKey(conversationId);
        redisChatTemplate.delete(key);
    }

    @NotNull
    private static String getConversationKey(String conversationId) {
        return "RCHAT_MEMORY_KEY_" + conversationId;
    }

    @Bean
    public RedisTemplate<String, List<Message>> redisChatTemplate(RedisConnectionFactory factory) {
        RedisTemplate<String, List<Message>> template = new RedisTemplate<>();
        template.setConnectionFactory(factory);

        template.setKeySerializer(new StringRedisSerializer());
        template.setHashKeySerializer(new StringRedisSerializer());

        KryoRedisSerializer serializer = new KryoRedisSerializer();
        template.setValueSerializer(serializer);
        template.setHashValueSerializer(serializer);

        template.afterPropertiesSet();
        return template;
    }

    public static class KryoRedisSerializer implements RedisSerializer<List<Message>> {
        private static final Kryo kryo = new Kryo();

        static {
            kryo.setRegistrationRequired(false);
            kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
        }

        @Override
        public byte[] serialize(List<Message> value) throws SerializationException {
            ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
            Output output = new Output(byteArrayOutputStream);

            kryo.writeObject(output, value);
            output.close();

            return byteArrayOutputStream.toByteArray();
        }

        @Override
        public List<Message> deserialize(byte[] bytes) throws SerializationException {
            if (bytes == null) {
                return new ArrayList<>();
            }
            Input input = new Input(bytes);
            List list = kryo.readObject(input, ArrayList.class);
            input.close();
            return list;
        }
    }
}
